import os
from datetime import datetime

import joblib
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import Ridge, SGDRegressor
from sklearn.metrics import r2_score, mean_squared_error
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import cross_val_score
import time
import warnings

warnings.filterwarnings("ignore", category=UserWarning)


class RidgeRegressionModel:
    def __init__(self,ID, task_type, model_type, excel_path, features, target, test_size=0.2, random_state=42, alpha=1.0):
        # 读取数据
        data = pd.read_excel(excel_path)
        self.df = pd.DataFrame(data)

        self.model_ID = ID
        self.task_type = task_type
        self.model_type = model_type

        # 初始化特征、目标
        self.X = self.df[features]
        self.y = self.df[target]
        self.features = features
        self.target = target

        # 数据分割
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(self.X, self.y,
                                                                                test_size=test_size,
                                                                                random_state=random_state)

        # 数据标准化
        self.scaler = StandardScaler()
        self.X_train_scaled = self.scaler.fit_transform(self.X_train)
        self.X_test_scaled = self.scaler.transform(self.X_test)

        # 初始化岭回归模型
        self.ridge_model = Ridge(alpha=alpha)

    def train(self):
        # 训练模型
        start_time_ridge = time.time()
        self.ridge_model.fit(self.X_train_scaled, self.y_train)
        self.ridge_train_time = time.time() - start_time_ridge

        # 保存训练好的模型
        current_time = datetime.now().strftime("%Y%m%d%H%M%S")
        # model_filename = f'../main/Data_Management_page_main/models/{current_time}_rid_model.pkl'  # 单独测试保存的路径
        # 创建一个model_ID文件夹用于存放模型
        if not os.path.exists(f"models/{self.model_ID}"):
            os.makedirs(f"models/{self.model_ID}")
        model_filename = f'./models/{self.model_ID}/{current_time}_rid_model.pkl'  # 整体测试保存的路径
        model_filename_result = model_filename[model_filename.rfind('/') + 1:model_filename.rfind('.')]

        print(model_filename_result)
        joblib.dump(self.ridge_model, model_filename)
        print(f"模型已保存到 {model_filename}")

        # 预测
        self.y_train_pred = self.ridge_model.predict(self.X_train_scaled)
        self.y_test_pred = self.ridge_model.predict(self.X_test_scaled)

        # 计算R² 和 MSE分数
        self.r2_train = r2_score(self.y_train, self.y_train_pred)
        self.r2_test = r2_score(self.y_test, self.y_test_pred)
        self.mse_train = mean_squared_error(self.y_train, self.y_train_pred)
        self.mse_test = mean_squared_error(self.y_test, self.y_test_pred)

        # 输出结果
        print(f"岭回归训练时间: {self.ridge_train_time:.2f} 秒")
        print(f"岭回归模型训练集R²: {self.r2_train}")
        print(f"岭回归模型测试集R²: {self.r2_test}")
        print(f"岭回归模型训练集MSE: {self.mse_train}")
        print(f"岭回归模型测试集MSE: {self.mse_test}")

        # 泛化能力（交叉验证得分）
        cv_scores = cross_val_score(self.ridge_model, self.X_train_scaled, self.y_train, cv=5, scoring='r2')
        print(f"岭回归交叉验证得分: {np.round(cv_scores, 4)}")
        cv_mean_score = np.mean(cv_scores)
        print(f"岭回归交叉验证平均得分: {cv_mean_score}")

        results = {
            "task_type": self.task_type,
            "model_type": self.model_type,
            "model_name": "岭回归",
            "train_time": self.ridge_train_time,
            "r2_train": self.r2_train,
            "r2_test": self.r2_test,
            "mse_train": self.mse_train,
            "mse_test": self.mse_test,
            "cv_scores": cv_scores.tolist(),
            "cv_mean_score": cv_mean_score,
            "model_filename": model_filename_result,
            "features": self.features,
            "target": self.target,
        }

        return results

    def plot_loss_curve(self):
        # 设置不同的 alpha 值（正则化强度）
        alphas = np.logspace(-4, 4, 100)  # 从 10^-4 到 10^4 的对数空间生成 100 个 alpha

        train_loss = []
        test_loss = []

        # 针对每个 alpha 值训练岭回归模型，并计算训练集和测试集的损失
        for alpha in alphas:
            ridge_model = Ridge(alpha=alpha)
            ridge_model.fit(self.X_train_scaled, self.y_train)

            # 计算训练集和测试集的均方误差
            y_train_pred = ridge_model.predict(self.X_train_scaled)
            y_test_pred = ridge_model.predict(self.X_test_scaled)

            train_loss.append(mean_squared_error(self.y_train, y_train_pred))
            test_loss.append(mean_squared_error(self.y_test, y_test_pred))

        # 绘制损失曲线
        plt.figure(figsize=(10, 6))
        plt.plot(alphas, train_loss, label='Training Loss', color='blue')
        plt.plot(alphas, test_loss, label='Testing Loss', color='orange')
        plt.xscale('log')  # alpha 取对数尺度
        plt.title('Ridge Regression Loss Curve (MSE vs Alpha)')
        plt.xlabel('Alpha (log scale)')
        plt.ylabel('Mean Squared Error')
        plt.legend()
        plt.grid(True)
        plt.savefig(r'C:\Users\13945\Desktop\picture\plot2.png')
        plt.show()
        plt.close()

# 创建主函数
def main(
        task_type=None,
        model_type=None,
        test_size=None,
        random_state=None,
        alpha=None):
    print(f"Parameters:{task_type}, {model_type}, {test_size}, {random_state}, {alpha}")

    # 数据集路径
    excel_path = r'C:\Users\13945\Desktop\MLDS铝合金成分设计数据.xlsx'
    features = ['Si/％', 'Mn/％', 'Zn/％', 'Mg/％', 'Cu/％', 'Cr/％', 'Zr/％', 'Ti/％', 'Fe/％', 'Ni/％', 'other']
    target = 'Ultimate tensile strength/MPa'

    # 创建模型实例
    model = RidgeRegressionModel(
        task_type=task_type,
        model_type=model_type,
        excel_path=excel_path,
        features=features,
        target=target,
        **{k: v for k, v in {
            'test_size': test_size,
            'random_state': random_state,
            'alpha': alpha
        }.items() if v is not None}  # 过滤掉None值的参数
    )

    # 训练模型并打印结果
    results = model.train()

    # 绘制图表
    model.plot_loss_curve()

    return results


if __name__ == "__main__":
    main()
